package simpledb.algorithm.join;

import simpledb.execution.JoinPredicate;
import simpledb.execution.OpIterator;
import simpledb.execution.Predicate;
import simpledb.storage.Tuple;
import simpledb.storage.TupleDesc;
import simpledb.storage.TupleIterator;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 排序归并连接（Sort-Merge Join）实现。
 *
 * 实现步骤如下：
 * 1. 将外表（左表）划分成多个块；
 * 2. 对每个左表块，将内表（右表）也划分成多个块；
 * 3. 分别对左右块内的元组进行排序；
 * 4. 使用归并方式在两个已排序的块之间执行连接操作。
 */
public class SortMergeJoin extends JoinStrategy {

    // 块缓存大小，默认为 131072 * 5 字节
    private final int     blockCacheSize = 131072 * 5;

    // 左右数据块（用于分块处理）
    private Tuple[]       block1;                     // 左表块
    private Tuple[]       block2;                     // 右表块

    // 比较谓词：小于、等于
    private JoinPredicate lt;
    private JoinPredicate eq;

    /**
     * 构造函数
     *
     * @param child1 左侧输入操作符（左表）
     * @param child2 右侧输入操作符（右表）
     * @param td 连接后的元组描述信息
     * @param joinPredicate 连接条件
     */
    public SortMergeJoin(final OpIterator child1, final OpIterator child2, final TupleDesc td,
                         final JoinPredicate joinPredicate) {
        super(child1, child2, td, joinPredicate);

        // 根据缓存大小计算每块能容纳的元组数量
        final int tuple1Num = this.blockCacheSize / child1.getTupleDesc().getSize();
        final int tuple2Num = this.blockCacheSize / child2.getTupleDesc().getSize();

        // 初始化两个内存块
        this.block1 = new Tuple[tuple1Num];
        this.block2 = new Tuple[tuple2Num];

        // 获取连接字段索引
        final int field1 = joinPredicate.getField1();
        final int field2 = joinPredicate.getField2();

        // 创建比较谓词
        this.lt = new JoinPredicate(field1, Predicate.Op.LESS_THAN, field2);
        this.eq = new JoinPredicate(field1, Predicate.Op.EQUALS, field2);
    }

    @Override
    public TupleIterator doJoin() {
        final List<Tuple> tupleList = new ArrayList<>();

        try {
            // 重置左表迭代器
            child1.rewind();
            while (child1.hasNext()) {
                // 从左表中读取一个块的数据
                int end1 = fetchTuples(child1, block1);

                // 重置右表迭代器
                child2.rewind();
                while (child2.hasNext()) {
                    // 从右表中读取一个块的数据
                    int end2 = fetchTuples(child2, block2);

                    // 对这两个块执行归并连接
                    mergeJoin(tupleList, end1, end2);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
            System.out.println("排序归并连接时发生错误: " + e.getMessage());
        }

        // 清空块数据
        Arrays.fill(this.block1, null);
        Arrays.fill(this.block2, null);

        // 返回结果元组的迭代器
        return new TupleIterator(this.td, tupleList);
    }

    /**
     * 执行两个块之间的归并连接操作
     *
     * @param tupleList 存储连接结果的列表
     * @param end1 block1 中实际使用的元组数
     * @param end2 block2 中实际使用的元组数
     */
    private void mergeJoin(final List<Tuple> tupleList, int end1, int end2) {
        // 1. 对两个块中的元组按连接字段排序
        final int field1 = this.joinPredicate.getField1();
        final int field2 = this.joinPredicate.getField2();
        sortTuples(block1, end1, field1); // 按左表字段排序
        sortTuples(block2, end2, field2); // 按右表字段排序

        // 2. 根据不同的连接类型执行归并连接
        int index1 = 0, index2 = 0;
        final Predicate.Op op = this.joinPredicate.getOperator();

        switch (op) {
            case EQUALS:
                // 等值连接
                while (index1 < end1 && index2 < end2) {
                    final Tuple lTuple = this.block1[index1];
                    final Tuple rTuple = this.block2[index2];

                    if (eq.filter(lTuple, rTuple)) {
                        // 找出所有与当前元组相等的元组范围
                        final JoinPredicate eq1 = new JoinPredicate(field1, Predicate.Op.EQUALS, field1);
                        final JoinPredicate eq2 = new JoinPredicate(field2, Predicate.Op.EQUALS, field2);
                        int begin1 = index1 + 1, begin2 = index2 + 1;

                        while (begin1 < end1 && eq1.filter(lTuple, this.block1[begin1]))
                            begin1++;
                        while (begin2 < end2 && eq2.filter(rTuple, this.block2[begin2]))
                            begin2++;

                        // 组合匹配的元组
                        for (int i = index1; i < begin1; i++) {
                            for (int j = index2; j < begin2; j++) {
                                tupleList.add(mergeTuple(this.block1[i], this.block2[j], this.td));
                            }
                        }

                        index1 = begin1;
                        index2 = begin2;
                    } else if (lt.filter(lTuple, rTuple)) {
                        index1++;
                    } else {
                        index2++;
                    }
                }
                break;

            case LESS_THAN:
            case LESS_THAN_OR_EQ:
                // 小于或小于等于连接
                while (index1 < end1) {
                    final Tuple lTuple = this.block1[index1++];
                    while (index2 < end2 && !this.joinPredicate.filter(lTuple, this.block2[index2]))
                        index2++;
                    while (index2 < end2) {
                        final Tuple rTuple = this.block2[index2++];
                        tupleList.add(mergeTuple(lTuple, rTuple, this.td));
                    }
                }
                break;

            case GREATER_THAN:
            case GREATER_THAN_OR_EQ:
                // 大于或大于等于连接
                while (index1 < end1) {
                    final Tuple lTuple = this.block1[index1++];
                    while (index2 < end2 && this.joinPredicate.filter(lTuple, this.block2[index2]))
                        index2++;
                    for (int i = 0; i < index2; i++) {
                        final Tuple rTuple = this.block2[i];
                        tupleList.add(mergeTuple(lTuple, rTuple, this.td));
                    }
                }
                break;
        }
    }

    /**
     * 对指定的元组数组按给定字段进行排序
     *
     * @param tuples 元组数组
     * @param len 数组有效长度
     * @param field 排序依据的字段索引
     */
    private void sortTuples(final Tuple[] tuples, int len, int field) {
        final JoinPredicate lt = new JoinPredicate(field, Predicate.Op.LESS_THAN, field);
        final JoinPredicate gt = new JoinPredicate(field, Predicate.Op.GREATER_THAN, field);
        Arrays.sort(tuples, 0, len, (o1, o2) -> {
            if (lt.filter(o1, o2)) {
                return -1;
            }
            if (gt.filter(o1, o2)) {
                return 1;
            }
            return 0;
        });
    }

    /**
     * 关闭连接操作，释放资源
     */
    @Override
    public void close() {
        this.block1 = null;
        this.block2 = null;
    }
}
